import os
import traceback
from tkinter import filedialog, messagebox
# 导入ruamel.yaml以保留注释
from ruamel.yaml import YAML

# 全局变量存储train.yaml路径和内容
_train_yaml_path = "../ultralytics/cfg/train.yaml"
_train_yaml_data = {}

# 初始化ruamel.yaml对象
_yaml = YAML()
_yaml.preserve_quotes = True
_yaml.indent(mapping=2, sequence=4, offset=2)


def init_train_yaml(fixed_path=None):
    """初始化train.yaml配置，设置固定路径"""
    global _train_yaml_path, _train_yaml_data

    if fixed_path:
        _train_yaml_path = os.path.abspath(fixed_path)
    else:
        # 默认路径：项目根目录下的train.yaml
        _train_yaml_path = os.path.abspath(os.path.join(
            os.path.dirname(__file__), "..", "train.yaml"
        ))

    return load_train_yaml()


def load_train_yaml():
    """加载固定路径的train.yaml文件内容"""
    global _train_yaml_data, _train_yaml_path

    if not _train_yaml_path:
        messagebox.showerror("未初始化", "请先调用init_train_yaml初始化配置")
        return False

    try:
        if not os.path.exists(_train_yaml_path):
            messagebox.showerror("文件不存在", f"train.yaml不存在: {_train_yaml_path}")
            return False

        if not os.access(_train_yaml_path, os.R_OK):
            messagebox.showerror("权限不足", f"无法读取train.yaml: {_train_yaml_path}")
            return False

        # 使用ruamel.yaml加载以保留注释
        with open(_train_yaml_path, 'r', encoding='utf-8') as f:
            _train_yaml_data = _yaml.load(f) or {}

        return True

    except Exception as e:
        messagebox.showerror("加载失败",
                             f"加载train.yaml出错:\n{str(e)}\n{traceback.format_exc()}")
        return False


def cmd_load_data(shared_data, app):
    """
    打开文件选择对话框选择data.yaml，
    并将其路径更新到train.yaml的'data'字段中
    """
    # 确保train.yaml已初始化
    if not _train_yaml_path:
        # 尝试使用默认路径初始化
        if not init_train_yaml():
            app.show_status("初始化train.yaml失败，无法加载数据配置")
            return False

    # 打开文件选择对话框，选择data.yaml
    file_path = filedialog.askopenfilename(
        title="选择data.yaml文件",
        filetypes=[("YAML文件", "*.yaml"), ("所有文件", "*.*")],
        defaultextension=".yaml"
    )

    if not file_path:
        app.show_status("已取消选择data.yaml文件")
        return False

    # 转换为绝对路径
    abs_data_path = os.path.abspath(file_path)

    # 验证文件是否存在
    if not os.path.exists(abs_data_path):
        messagebox.showerror("文件不存在", f"所选data.yaml不存在: {abs_data_path}")
        return False

    # 更新到train.yaml中
    result = update_train_data_fields(
        fields_to_update={'data': abs_data_path},
        app=app
    )

    if result:
        app.show_status(f"已成功加载数据配置: {os.path.basename(abs_data_path)}")
        # 显示成功弹窗
        messagebox.showinfo("加载成功", f"已成功加载数据配置:\n{abs_data_path}")
        # 同时更新共享数据
        if hasattr(shared_data, 'data_config'):
            shared_data.data_config['data_path'] = abs_data_path
    return result


def update_train_data_fields(fields_to_update, app=None):
    """更新train.yaml中的指定字段"""
    global _train_yaml_data, _train_yaml_path

    if not _train_yaml_data:
        if not load_train_yaml():
            return False

    if not fields_to_update:
        _show_status(app, "没有需要更新的数据字段")
        return True

    try:
        # 重新加载配置以确保获取最新的带注释内容
        with open(_train_yaml_path, 'r', encoding='utf-8') as f:
            config = _yaml.load(f) or {}

        for field, value in fields_to_update.items():
            # 处理路径字段
            if field in ['data', 'train', 'val'] and value and not os.path.isabs(value):
                value = os.path.abspath(os.path.join(
                    os.path.dirname(_train_yaml_path),
                    value
                ))

            # 验证路径
            if field in ['data', 'train', 'val'] and value and os.path.isfile(value) and not os.path.exists(value):
                messagebox.showwarning("路径无效", f"文件不存在: {value}")
                continue

            config[field] = value
            _show_status(app, f"已更新字段: {field} = {value}")

        # 检查写入权限
        if not os.access(_train_yaml_path, os.W_OK):
            messagebox.showerror("权限不足", f"无法写入train.yaml: {_train_yaml_path}")
            return False

        # 使用ruamel.yaml写回文件（保留注释）
        with open(_train_yaml_path, 'w', encoding='utf-8') as f:
            _yaml.dump(config, f)

        _show_status(app, f"train.yaml已更新: {_train_yaml_path}")
        return True

    except Exception as e:
        _show_status(app, f"更新失败: {str(e)}")
        messagebox.showerror("更新失败",
                             f"更新train.yaml出错:\n{str(e)}\n{traceback.format_exc()}")
        return False


def get_train_data_field(field_name):
    """获取train.yaml中的指定字段值"""
    global _train_yaml_data
    return _train_yaml_data.get(field_name)


def _show_status(app, message):
    """显示状态信息"""
    if app and hasattr(app, 'show_status'):
        app.show_status(message)
    else:
        print(f"[数据操作] {message}")